import csv
import os
import pickle

import logger
from utils.file import exist_create_dir

log = logger.get_logger()


def write_pkl(path, dataSet):
    """
    持久化数据集
    :param path: 持久化目录
    :param dataSet: 数据集
    :return:
    """
    saveAbsPath = os.path.abspath(path)
    exist_create_dir(saveAbsPath)
    log.info("save pkl int this path : " + saveAbsPath)
    with open(saveAbsPath, "wb") as file:
        pickle.dump(dataSet, file)


def write_csv(path, data):
    """
    持久化
    :param path: 存储路径
    :param data:
    :return:
    """
    saveAbsPath = os.path.abspath(path)
    exist_create_dir(saveAbsPath)
    log.info("save csv int this path : " + saveAbsPath)
    with open(saveAbsPath, "w") as file:
        writer = csv.writer(file)
        writer.writerow(data)


def write_txt(path, data):
    """
    写入txt文件
    :param path:
    :param data:
    :return:
    """
    saveAbsPath = os.path.abspath(path)
    exist_create_dir(saveAbsPath)
    log.info("save str int this path : " + saveAbsPath)
    with open(saveAbsPath, "w") as file:
        file.write(str(data))
